package transformers;

import asms.LoaderVisitor;
import asms.ProcessBuilderVisitor;
import org.objectweb.asm.*;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class AgentTransformer implements ClassFileTransformer {

    /**
     * 加载的每一个类都会进入到transform方法，这里我们筛选一下，只处理指定的controller.Load类。
     * */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replace("/",".");
        boolean flag = false;
        if(className.equals("java.lang.ProcessBuilder")){
            System.out.println("I find the "+ className + ",then I'll try to hook this class.");
            ClassReader classReader = new ClassReader(classfileBuffer);
            ClassWriter classWriter = new ClassWriter(classReader, Opcodes.ASM5);
            ClassVisitor classVisitor = new ProcessBuilderVisitor(classWriter);
            classReader.accept(classVisitor,ClassReader.EXPAND_FRAMES);
            byte[] bytes = classWriter.toByteArray();

            if (flag) {
                // 调试的时候如果想看修改后类的字节码，可以将flag设置为true
                try {
                    FileOutputStream fos = new FileOutputStream(new File("/tmp/Target.class"));
                    fos.write(bytes);
                    fos.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }

            return bytes;
        }

        return classfileBuffer;
    }
}